from UItools import vtk
import numpy as np

class CustomInteractorStyle(vtk.vtkInteractorStyleTrackballCamera):
    def __init__(self, parent=None):
        self.AddObserver("LeftButtonPressEvent", self.left_button_press_event)
        self.AddObserver("LeftButtonReleaseEvent", self.left_button_release_event)
        self.AddObserver("MiddleButtonPressEvent", self.middle_button_press_event)
        self.AddObserver("MiddleButtonReleaseEvent", self.middle_button_release_event)
        self.AddObserver("MouseMoveEvent", self.mouse_move_event)
        self.AddObserver("MouseWheelForwardEvent", self.mouse_wheel_forward_event)
        self.AddObserver("MouseWheelBackwardEvent", self.mouse_wheel_backward_event)
        
        self.left_button_pressed = False
        self.middle_button_pressed = False
        self.shift_pressed = False
        self.last_pos = None

    def left_button_press_event(self, obj, event):
        self.left_button_pressed = True
        self.shift_pressed = self.GetInteractor().GetShiftKey()
        self.last_pos = self.GetInteractor().GetEventPosition()
        self.OnLeftButtonDown()  
        return

    def left_button_release_event(self, obj, event):
        self.left_button_pressed = False
        self.shift_pressed = False
        self.OnLeftButtonUp()
        return

    def middle_button_press_event(self, obj, event):
        self.middle_button_pressed = True
        self.last_pos = self.GetInteractor().GetEventPosition()
        self.OnMiddleButtonDown()  
        return

    def middle_button_release_event(self, obj, event):
        self.middle_button_pressed = False
        self.OnMiddleButtonUp()
        return

    def mouse_move_event(self, obj, event):
        current_pos = self.GetInteractor().GetEventPosition()
        if not self.last_pos:
            self.last_pos = current_pos
            return
            
        dx = current_pos[0] - self.last_pos[0]
        dy = current_pos[1] - self.last_pos[1]
        
        renderer = self.GetCurrentRenderer()
        if not renderer:
            return
            
        camera = renderer.GetActiveCamera()

        # 中键拖动：旋转
        if self.middle_button_pressed and not self.shift_pressed:
            camera.Azimuth(-dx * 0.5)
            camera.Elevation(-dy * 0.5)
            camera.OrthogonalizeViewUp()
        
        # Shift+左键：平移
        elif self.left_button_pressed and self.shift_pressed:
            # 计算正确的平移方向
            focal_point = np.array(camera.GetFocalPoint())
            position = np.array(camera.GetPosition())
            view_up = np.array(camera.GetViewUp())
            
            view_direction = focal_point - position
            distance = np.linalg.norm(view_direction)
            view_direction = view_direction / distance
            
            right = np.cross(view_up, view_direction)
            right = right / np.linalg.norm(right)
            
            # 计算平移向量
            translation = (dx * right - dy * view_up) * distance * 0.001
            
            new_position = position + translation
            new_focal_point = focal_point + translation
            
            camera.SetPosition(new_position)
            camera.SetFocalPoint(new_focal_point)
        
        self.last_pos = current_pos
        self.GetInteractor().Render()
        return

    def mouse_wheel_forward_event(self, obj, event):
        renderer = self.GetCurrentRenderer()
        if renderer:
            camera = renderer.GetActiveCamera()
            camera.Dolly(1.1)
            renderer.ResetCameraClippingRange()
            self.GetInteractor().Render()
        return

    def mouse_wheel_backward_event(self, obj, event):
        renderer = self.GetCurrentRenderer()
        if renderer:
            camera = renderer.GetActiveCamera()
            camera.Dolly(0.9)
            renderer.ResetCameraClippingRange()
            self.GetInteractor().Render()
        return